# GoogleSearch class, gets search results from Google
# Can be treated like a list
# @author Richard Cornelissen

import math, urllib, re, sys
from HIT.browser import Browser, BrowserError
from HIT.metasearch import SearchEngine

class GoogleSearch(SearchEngine):
    
    GOOGLE_URL = "http://m.google.com/m?safe=off&gwt=off&hl=en&q=%(query)s&start=%(start)d"
    NAME = u'Google'
    LOGO = u'google.png'
    MAX_RESULTS = 1000
    
    def __init__(self, query, cache_before = 20):
        self._results = []
        self.query = unicode(query)
        self._get_results(cache_before)
    
    # __getitem__ function, enables the use of GoogleSearch[x]
    def __getitem__(self, key):
        if isinstance(key, int):
            try:
                if (key >= self.MAX_RESULTS):
                    return None
                
                if (key >= len(self._results)):
                    self._get_results(key)
            
                return self._results[key]
            except BrowserError:
                return None
            
        elif isinstance(key, slice):
            try:
                if key.start > self.MAX_RESULTS:
                    return [None]
                
                if key.stop > self.MAX_RESULTS:
                    key.stop = self.MAX_RESULTS
                
                if (len(self._results) <= key.stop):
                    self._get_results(key.stop)
                
                return self._results[key]
            except BrowserError:
                return [None]
    
    def __len__(self):
        return len(self._results)
    
    # Function that gets results until max results are cached
    def _get_results(self, max):
        max = math.ceil(max/10.0) * 10
        
        while len(self._results) <= max:
            start = len(self._results)
            soup = self._get_results_page(start=start)
            newresults = self._extract_results(soup)
            
            self._results = self._results + newresults
            
        return self._results
    
    # Function that gets a page from Google
    def _get_results_page(self, start = 0):
        url = GoogleSearch.GOOGLE_URL % {   'query': urllib.quote_plus(self.query),
                                            'start': start 
                                        }
        b = Browser()
        page = b.get_page(url)
        from BeautifulSoup import BeautifulSoup
        return BeautifulSoup(page)
    
    # Function that extracts multiple results
    def _extract_results(self, soup):
        #results = soup.findAll('li', {'class': 'g'})
        results = soup.findAll('div', {'class': 'r ld'})
        ret_res = []
        for result in results:
            eres = self._extract_result(result)
            if eres:
                ret_res.append(eres)
        return ret_res
    
    # Function that extracts one result, used by _extract_result
    def _extract_result(self, result):
        title, url = self._extract_title_url(result)
        desc = self._extract_description(result)
        if not title or not url or not desc:
            return None
        return {'title': title, 'url': url, 'desc': desc} 
    
    def _extract_title_url(self, result):
        title_a = result.find('a')
        
        if not title_a:
            return None, None
        title = ''.join(title_a.findAll(text=True))
        title = self._html_unescape(title)
        url = title_a['href']
        
        match = re.match(r'/m/url\?(.)*q=(http[^&]+)&(.)*', url)
        if match:
            url = urllib.unquote(match.group(2))
        return title, url
    
    def _extract_description(self, result):
        desc_div = result.find('div', {'class': 'kd'})
        if not desc_div:
            return None
        
        desc = ''.join(desc_div.findAll(text=True))
        return self._html_unescape(desc)
    
    # Unescaped a HTML string
    def _html_unescape(self, str):
        def entity_replacer(m):
            entity = m.group(1)
            from htmlentitydefs import name2codepoint
            if entity in name2codepoint:
                return unichr(name2codepoint[entity])
            else:
                return m.group(0)
        
        def ascii_replacer(m):
            cp = int(m.group(1))
            if cp <= 255:
                return unichr(cp)
            else:
                return m.group(0)
        
        s =    re.sub(r'&#(\d+);',  ascii_replacer, str, re.U)
        return re.sub(r'&([^;]+);', entity_replacer, s, re.U)